from typing import List

from llama_index.core.agent.workflow import  FunctionAgent
from llama_index.core.indices.common.struct_store.sql import SQLStructDatapointExtractor
from llama_index.core.indices.struct_store import NLSQLTableQueryEngine
from llama_index.core.node_parser import SimpleNodeParser
from llama_index.core.query_engine.flare.answer_inserter import LLMLookaheadAnswerInserter
from llama_index.core.tools import QueryEngineTool
from llama_index.core.vector_stores import SimpleVectorStore
from llama_index.core.schema import  TextNode
from llama_index.core import Settings, SimpleKeywordTableIndex, SummaryIndex, get_response_synthesizer, \
    DocumentSummaryIndex, SimpleDirectoryReader, VectorStoreIndex
from llama_index.embeddings.zhipuai import ZhipuAIEmbedding
from llama_index.core.graph_stores import SimplePropertyGraphStore
from llama_index.core.schema import Document
from pydantic import BaseModel
from llama_index.core.indices.property_graph.base import PropertyGraphIndex
from llama_index.core.indices.property_graph.retriever import PGRetriever
from llama_index.core.indices.property_graph.sub_retrievers.base import BasePGRetriever
from llama_index.core.indices.property_graph.sub_retrievers.custom import (
    CustomPGRetriever,
    CUSTOM_RETRIEVE_TYPE,
)
from llama_index.core.indices.property_graph.sub_retrievers.cypher_template import (
    CypherTemplateRetriever,
)
from llama_index.core.indices.property_graph.sub_retrievers.llm_synonym import (
    LLMSynonymRetriever,
)
from llama_index.core.indices.property_graph.sub_retrievers.text_to_cypher import (
    TextToCypherRetriever,
)
from llama_index.core.indices.property_graph.sub_retrievers.vector import (
    VectorContextRetriever,
)
from llama_index.core.indices.property_graph.transformations.implicit import (
    ImplicitPathExtractor,
)
from llama_index.core.indices.property_graph.transformations.schema_llm import (
    SchemaLLMPathExtractor,
)
from llama_index.core.indices.property_graph.transformations.simple_llm import (
    SimpleLLMPathExtractor,
)
from llama_index.core.indices.property_graph.transformations.dynamic_llm import (
    DynamicLLMPathExtractor,
)
from llama_index.core.indices.property_graph.utils import default_parse_triplets_fn

embed_model = ZhipuAIEmbedding(
    model="embedding-2",
    api_key="f387f5e4837d4e4bba6d267682a957c9.PmPiTw8qVlsI2Oi5"
    # With the `embedding-3` class
    # of models, you can specify the size
    # of the embeddings you want returned.
    # dimensions=1024
)
Settings.embed_model=embed_model

from llama_index.llms.deepseek import DeepSeek

llm = DeepSeek(model="deepseek-chat", api_key="sk-605e60a1301040759a821b6b677556fb")
Settings.llm = llm

from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, insert

engine = create_engine("sqlite:///:memory:", future=True)
metadata = MetaData()
city_stats = Table(
    "city_stats", metadata,
    Column("city_name", String(16), primary_key=True),
    Column("population", Integer),
    Column("country", String(16), nullable=False)
)
metadata.create_all(engine)

data = [
    {"city_name": "Toronto", "population": 2731571, "country": "Canada"},
    {"city_name": "Tokyo", "population": 13929286, "country": "Japan"}
]
with engine.begin() as conn:
    for row in data:
        conn.execute(city_stats.insert().values(**row))


from llama_index.core import SQLDatabase, VectorStoreIndex
from llama_index.core.query_engine import SQLJoinQueryEngine

# 包装SQLAlchemy引擎
sql_database = SQLDatabase(engine)

documents=[Document(text='''**东京经济报告摘要：**
东京作为日本的首都和全球重要经济中心，其经济特点包括：
- 全球重要的金融中心，拥有东京证券交易所
- 高度发达的服务业和制造业
- 科技创新和研发中心
- 高度集中的企业总部和跨国公司
- 发达的交通基础设施和物流网络''')]
# 假设已存在向量存储索引（需预先加载文档）
vector_index = VectorStoreIndex.from_documents(documents)

queryEngine =NLSQLTableQueryEngine(sql_database)

queryEngineTool=QueryEngineTool.from_defaults(queryEngine,description="城市信息")

other_query_tool=QueryEngineTool.from_defaults(vector_index.as_query_engine(),description="东京经济报告摘要")
# 创建联合查询引擎
query_engine = SQLJoinQueryEngine(
    sql_query_tool=queryEngineTool,
    other_query_tool=other_query_tool
)

response = query_engine.query(
    "找出人口超过500万的城市，并列出相关经济报告摘要"
)
print(response)